// SPDX-License-Identifier: Apache-2.0
// Copyright 2021 Authors of KubeArmor

package anonmapexec

import (
	"bytes"
	"encoding/binary"
	"errors"
	"log"
	"os"
	"strings"
	"sync"

	"github.com/cilium/ebpf"
	"github.com/cilium/ebpf/link"
	"github.com/cilium/ebpf/ringbuf"
	"github.com/cilium/ebpf/rlimit"
	"github.com/kubearmor/KubeArmor/KubeArmor/buildinfo"
	"github.com/kubearmor/KubeArmor/KubeArmor/presets/base"

	fd "github.com/kubearmor/KubeArmor/KubeArmor/feeder"
	mon "github.com/kubearmor/KubeArmor/KubeArmor/monitor"
	tp "github.com/kubearmor/KubeArmor/KubeArmor/types"
)

//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -cc clang anonmapexec  ../../BPF/anonmapexec.bpf.c -type mmap_event -no-global-types -- -I/usr/include/ -O2 -g

var (
	_ base.PresetInterface = (*Preset)(nil)
)

const (
	NAME   string = "AnonMapExecutionPreset"
	EXISTS uint32 = 1024
)

type NsKey struct {
	PidNS uint32
	MntNS uint32
}

type ContainerVal struct {
	NsKey  NsKey
	Policy string
}

type Preset struct {
	base.Preset

	BPFContainerMap *ebpf.Map

	// events
	Events        *ringbuf.Reader
	EventsChannel chan []byte

	// ContainerID -> NsKey
	ContainerMap     map[string]ContainerVal
	ContainerMapLock *sync.RWMutex

	Link link.Link

	obj anonmapexecObjects
}

type MmapEvent struct {
	Ts uint64

	PidID uint32
	MntID uint32

	HostPPID uint32
	HostPID  uint32

	PPID uint32
	PID  uint32
	UID  uint32

	EventID int32

	Retval int64

	Comm [80]byte

	TTY [64]byte

	Args [6]uint64
}

func NewAnonMapExecPreset() *Preset {
	p := &Preset{}
	p.ContainerMap = make(map[string]ContainerVal)
	p.ContainerMapLock = new(sync.RWMutex)
	return p
}

func (p *Preset) Name() string {
	return NAME
}

func (p *Preset) RegisterPreset(logger *fd.Feeder, monitor *mon.SystemMonitor) (base.PresetInterface, error) {

	if logger.Enforcer != "BPFLSM" {
		// it's based on active enforcer, it might possible that node support bpflsm but
		// current enforcer is not bpflsm
		return nil, errors.New("AnonExecutionPreset not supported if bpflsm not supported")
	}

	p.Logger = logger
	p.Monitor = monitor
	var err error

	if err = rlimit.RemoveMemlock(); err != nil {
		p.Logger.Errf("Error removing rlimit %v", err)
		return nil, err // Doesn't require clean up so not returning err
	}

	p.BPFContainerMap, _ = ebpf.NewMapWithOptions(&ebpf.MapSpec{
		Type:       ebpf.Hash,
		KeySize:    8,
		ValueSize:  4,
		MaxEntries: 256,
		Pinning:    ebpf.PinByName,
		Name:       "kubearmor_anon_map_exec_preset_containers",
	}, ebpf.MapOptions{
		PinPath: monitor.PinPath,
	})

	if err := loadAnonmapexecObjects(&p.obj, &ebpf.CollectionOptions{
		Maps: ebpf.MapOptions{
			PinPath: monitor.PinPath,
		},
	}); err != nil {
		p.Logger.Errf("Error loading BPF LSM objects: %v", err)
		return nil, err
	}

	p.Link, err = link.AttachLSM(link.LSMOptions{Program: p.obj.EnforceMmapFile})
	if err != nil {
		p.Logger.Errf("opening lsm %s: %s", p.obj.EnforceMmapFile.String(), err)
		return nil, err
	}

	p.Events, err = ringbuf.NewReader(p.obj.Events)
	if err != nil {
		p.Logger.Errf("opening ringbuf reader: %s", err)
		return nil, err
	}
	p.EventsChannel = make(chan []byte, mon.SyscallChannelSize)

	go p.TraceEvents()

	return p, nil

}

// TraceEvents traces events generated by bpflsm enforcer
func (p *Preset) TraceEvents() {

	if p.Events == nil {
		p.Logger.Err("ringbuf reader is nil, exiting trace events")
	}
	p.Logger.Print("Starting TraceEvents from AnonMapExec Presets")
	go func() {
		for {

			record, err := p.Events.Read()
			if err != nil {
				if errors.Is(err, ringbuf.ErrClosed) {
					// This should only happen when we call DestroyMonitor while terminating the process.
					// Adding a Warn just in case it happens at runtime, to help debug
					p.Logger.Warnf("Ring Buffer closed, exiting TraceEvents %s", err.Error())
					return
				}
				p.Logger.Warnf("Ringbuf error reading %s", err.Error())
				continue
			}

			p.EventsChannel <- record.RawSample

		}
	}()

	for {

		dataRaw := <-p.EventsChannel

		var event MmapEvent

		if err := binary.Read(bytes.NewBuffer(dataRaw), binary.LittleEndian, &event); err != nil {
			log.Printf("parsing ringbuf event: %s", err)
			continue
		}

		readLink := true

		containerID := ""

		if event.PidID != 0 && event.MntID != 0 {
			containerID = p.Monitor.LookupContainerID(event.PidID, event.MntID)
		}

		log := p.Monitor.BuildLogBase(event.EventID, mon.ContextCombined{
			ContainerID: containerID,
			ContextSys: mon.SyscallContext{
				PID:  event.PID,
				PPID: event.PPID,
				UID:  event.UID,

				HostPID:  event.HostPID,
				HostPPID: event.HostPPID,
				TTY:      event.TTY,
			},
		}, readLink)

		if ckv, ok := p.ContainerMap[containerID]; ok {
			log.PolicyName = ckv.Policy
			log.Type = "MatchedPolicy"
		}

		var f []string
		f = append(f, ParseProtectionFlags(event.Args[0]))
		f = append(f, ParseProtectionFlags(event.Args[1]))
		f = append(f, ParseMemoryFlags(event.Args[2]))
		if len(f) > 0 {
			log.Data = strings.Join(f, ",")
		}
		log.Operation = "File"
		if event.Retval >= 0 {
			log.Result = "Passed"
			log.Action = "Audit"
		} else {
			log.Result = "Permission denied"
			log.Action = "Block"
		}
		log.Enforcer = base.PRESET_ENFORCER + NAME
		log.KubeArmorVersion = buildinfo.GitSummary
		p.Logger.PushLog(log)

	}
}

func (p *Preset) RegisterContainer(containerID string, pidns, mntns uint32) {
	ckv := NsKey{PidNS: pidns, MntNS: mntns}

	p.ContainerMapLock.Lock()
	defer p.ContainerMapLock.Unlock()
	p.Logger.Debugf("[AnonMapExec] Registered container with id: %s\n", containerID)
	p.ContainerMap[containerID] = ContainerVal{NsKey: ckv}
}

func (p *Preset) UnregisterContainer(containerID string) {
	p.ContainerMapLock.Lock()
	defer p.ContainerMapLock.Unlock()

	if val, ok := p.ContainerMap[containerID]; ok {
		if err := p.DeleteContainerIDFromMap(containerID, val.NsKey); err != nil {
			p.Logger.Errf("Error deleting container %s: %s", containerID, err.Error())
			return
		}
		p.Logger.Debugf("[AnonMapExec] Unregistered container with id: %s\n", containerID)
		delete(p.ContainerMap, containerID)
	}
}

func (p *Preset) AddContainerIDToMap(id string, ckv NsKey, action string) error {
	p.Logger.Debugf("[AnonMapExec] adding container with id to anon_map exec map: %s\n", id)
	a := base.Block
	if action == "Audit" {
		a = base.Audit
	}
	if err := p.BPFContainerMap.Put(ckv, a); err != nil {
		p.Logger.Errf("Error adding container %s to outer map: %s", id, err)
		return err
	}
	return nil
}

func (p *Preset) DeleteContainerIDFromMap(id string, ckv NsKey) error {
	p.Logger.Debugf("[AnonMapExec] deleting container with id to anon_map exec map: %s\n", id)
	if err := p.BPFContainerMap.Delete(ckv); err != nil {
		if !errors.Is(err, os.ErrNotExist) {
			p.Logger.Errf("Error deleting container %s in kubearmor_anon_map_exec_preset_containers map: %s", id, err.Error())
			return err
		}
	}
	return nil
}

func (p *Preset) UpdateSecurityPolicies(endPoint tp.EndPoint) {
	var anonMapExecPresetRulePresent bool
	for _, cid := range endPoint.Containers {
		anonMapExecPresetRulePresent = false
		for _, secPolicy := range endPoint.SecurityPolicies {
			for _, preset := range secPolicy.Spec.Presets {
				if preset.Name == tp.AnonMapExec {
					anonMapExecPresetRulePresent = true
					p.ContainerMapLock.RLock()
					// Check if Container ID is registered in Map or not
					ckv, ok := p.ContainerMap[cid]
					if !ok {
						// It maybe possible that CRI has unregistered the containers but K8s construct still has not sent this update while the policy was being applied,
						// so the need to check if the container is present in the map before we apply policy.
						p.ContainerMapLock.RUnlock()
						return
					}
					ckv.Policy = secPolicy.Metadata["policyName"]
					p.ContainerMap[cid] = ckv
					err := p.AddContainerIDToMap(cid, ckv.NsKey, preset.Action)
					if err != nil {
						p.Logger.Warnf("Updating policy for container %s :%s ", cid, err)
					}
					p.ContainerMapLock.RUnlock()
				}
			}
		}
		if !anonMapExecPresetRulePresent {
			p.ContainerMapLock.RLock()
			ckv := p.ContainerMap[cid]
			_ = p.DeleteContainerIDFromMap(cid, ckv.NsKey)
			p.ContainerMapLock.RUnlock()
		}
	}
}

func (p *Preset) Destroy() error {
	if p == nil {
		return nil
	}
	var errBPFCleanUp error

	if err := p.obj.Close(); err != nil {
		p.Logger.Err(err.Error())
		errBPFCleanUp = errors.Join(errBPFCleanUp, err)
	}

	if err := p.Link.Close(); err != nil {
		p.Logger.Err(err.Error())
		errBPFCleanUp = errors.Join(errBPFCleanUp, err)
	}

	p.ContainerMapLock.Lock()

	if p.BPFContainerMap != nil {
		if err := p.BPFContainerMap.Unpin(); err != nil {
			p.Logger.Err(err.Error())
			errBPFCleanUp = errors.Join(errBPFCleanUp, err)
		}
		if err := p.BPFContainerMap.Close(); err != nil {
			p.Logger.Err(err.Error())
			errBPFCleanUp = errors.Join(errBPFCleanUp, err)
		}
	}

	p.ContainerMapLock.Unlock()

	if p.Events != nil {
		if err := p.Events.Close(); err != nil {
			p.Logger.Err(err.Error())
			errBPFCleanUp = errors.Join(errBPFCleanUp, err)
		}
	}

	p = nil
	return errBPFCleanUp
}
